{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Managing Multiple MCP Servers with MCPClientSessionManager"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "**Author:** Priyanshu Deshmukh\n",
    "\n",
    "This notebook demonstrates how to use the `MCPClientSessionManager` class to dynamically manage multiple MCP (Model Context Protocol) server sessions. In this example, we build a research assistant that can query both arXiv papers and Wikipedia articles by opening MCP sessions on-demand within a tool function.\n",
    "\n",
    "## Key Features Demonstrated\n",
    "\n",
    "- **Dynamic Session Management**: Open MCP sessions inside tool functions as needed\n",
    "- **Multi-Transport Support**: Handle both stdio (ArXiv) and SSE (Wikipedia) servers\n",
    "- **Tool-Based Server Selection**: Let agents choose which server to use\n",
    "- **Context Manager Pattern**: Automatic resource cleanup with `async with`\n",
    "- **Agent Handoffs**: Coordinate between research assistant and server-specific agents"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## What is MCPClientSessionManager?\n",
    "\n",
    "The `MCPClientSessionManager` is a utility class that simplifies managing MCP client sessions. Unlike the traditional approach where you open a session at the start and keep it alive, `MCPClientSessionManager` allows you to:\n",
    "\n",
    "- **Open sessions on-demand**: Create sessions only when needed within your workflow\n",
    "- **Switch between servers**: Dynamically select which MCP server to connect to\n",
    "- **Manage multiple transports**: Handle both `stdio` (process-based) and `SSE` (HTTP-based) protocols\n",
    "- **Ensure clean resource management**: Automatic cleanup with async context managers\n",
    "\n",
    "### Key Components\n",
    "\n",
    "1. **StdioConfig**: Configuration for stdio-based MCP servers (local processes)\n",
    "   - Starts a Python process that communicates via stdin/stdout\n",
    "   - Example: Local arXiv paper search server\n",
    "\n",
    "2. **SseConfig**: Configuration for SSE-based MCP servers (HTTP endpoints)\n",
    "   - Connects to a remote server via Server-Sent Events\n",
    "   - Example: Remote Wikipedia API server\n",
    "\n",
    "3. **MCPConfig**: Container for multiple server configurations\n",
    "   - Holds all available servers in one configuration object\n",
    "   - Enables dynamic server selection at runtime"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Setup and Imports\n",
    "\n",
    "Import all necessary libraries for our multi-server research assistant."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Only needed for Jupyter notebooks\n",
    "import nest_asyncio\n",
    "\n",
    "from autogen import ConversableAgent, LLMConfig\n",
    "from autogen.agentchat.group import AgentTarget\n",
    "from autogen.agentchat.group.llm_condition import StringLLMCondition\n",
    "from autogen.agentchat.group.on_condition import OnCondition\n",
    "from autogen.agentchat.group.reply_result import ReplyResult\n",
    "from autogen.mcp.mcp_client import MCPClientSessionManager, MCPConfig, SseConfig, StdioConfig, create_toolkit\n",
    "from autogen.tools import tool\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## Configure LLM\n",
    "\n",
    "Set up the language model configuration. Make sure to set your `OPENAI_API_KEY` environment variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "OPENAI_API_KEY = os.getenv(\"OPENAI_API_KEY\")\n",
    "\n",
    "llm_config = LLMConfig(\n",
    "    config_list=[\n",
    "        {\n",
    "            \"model\": \"gpt-4o\",\n",
    "            \"api_type\": \"openai\",\n",
    "            \"api_key\": OPENAI_API_KEY,\n",
    "        }\n",
    "    ],\n",
    "    temperature=0.7,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Configure Multiple MCP Servers\n",
    "\n",
    "Here we define two different MCP servers with different transport protocols:\n",
    "\n",
    "1. **ArxivServer** (stdio transport):\n",
    "   - Runs as a local Python process\n",
    "   - Provides tools to search and retrieve arXiv papers\n",
    "   - Stores papers in `/tmp/arxiv_papers`\n",
    "\n",
    "2. **WikipediaServer** (SSE transport):\n",
    "   - Connects to an HTTP endpoint\n",
    "   - Provides tools to search and retrieve Wikipedia articles\n",
    "   - Must be running separately (e.g., via `python mcp/mcp_wikipedia.py sse`)\n",
    "\n",
    "The `MCPConfig` object holds both server configurations, allowing us to dynamically select which one to use at runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Configure a stdio-based MCP server (local process)\n",
    "ArxivServer = StdioConfig(\n",
    "    command=\"python3\",\n",
    "    args=[\"mcp/mcp_arxiv.py\", \"stdio\", \"--storage-path\", \"/tmp/arxiv_papers\"],\n",
    "    transport=\"stdio\",\n",
    "    server_name=\"ArxivServer\",\n",
    ")\n",
    "\n",
    "# Configure an SSE-based MCP server (HTTP endpoint)\n",
    "WikipediaServer = SseConfig(\n",
    "    url=\"http://127.0.0.1:8000/sse\",\n",
    "    timeout=10,\n",
    "    sse_read_timeout=60,\n",
    "    server_name=\"WikipediaServer\",\n",
    ")\n",
    "\n",
    "# Create an MCPConfig with both servers\n",
    "mcp_config = MCPConfig(servers=[ArxivServer, WikipediaServer])\n",
    "\n",
    "print(f\"Configured {len(mcp_config.servers)} MCP servers:\")\n",
    "for server in mcp_config.servers:\n",
    "    print(f\"  - {server.server_name}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "9",
   "metadata": {},
   "source": [
    "## Helper Function: Get Server Config\n",
    "\n",
    "This utility function retrieves a specific server configuration by name from our `MCPConfig`. It's essential for dynamic server selection within tools."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_server_config(mcp_config: MCPConfig, server_name: str) -> StdioConfig | SseConfig:\n",
    "    \"\"\"\n",
    "    Return the server config (StdioConfig or SseConfig) matching the given server_name.\n",
    "    Args:\n",
    "        mcp_config: The MCP configuration containing all servers\n",
    "        server_name: Name of the server to retrieve\n",
    "    Returns:\n",
    "        The matching server configuration\n",
    "    Raises:\n",
    "        KeyError: If the server name is not found\n",
    "    \"\"\"\n",
    "    existing_names = {getattr(server, \"server_name\", None) for server in mcp_config.servers}\n",
    "    for server in mcp_config.servers:\n",
    "        if getattr(server, \"server_name\", None) == server_name:\n",
    "            return server\n",
    "    raise KeyError(f\"Server '{server_name}' not found in MCPConfig. Existing servers: {list(existing_names)}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "## Define the Research Assistant Agent\n",
    "\n",
    "The research assistant is the main agent that users interact with. It has access to a tool that can query either the ArXiv or Wikipedia servers based on the user's needs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "RESEARCH_AGENT_PROMPT = \"\"\"\n",
    "You are a research assistant agent.\n",
    "You will provide assistance for research tasks.\n",
    "You have two MCP servers to use:\n",
    "1. ArxivServer: to search for papers on arXiv\n",
    "2. WikipediaServer: to search for articles on Wikipedia\n",
    "\"\"\"\n",
    "\n",
    "research_assistant = ConversableAgent(\n",
    "    name=\"research_assistant\",\n",
    "    system_message=RESEARCH_AGENT_PROMPT,\n",
    "    llm_config=llm_config,\n",
    "    human_input_mode=\"NEVER\",\n",
    ")\n",
    "\n",
    "print(\"Research assistant created successfully!\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "13",
   "metadata": {},
   "source": [
    "## Create the MCP Tool: Dynamic Server Connection\n",
    "\n",
    "This is the key innovation: a tool that opens an MCP session on-demand, creates a specialized agent with the server's tools, executes the query, and returns results.\n",
    "\n",
    "### How It Works\n",
    "\n",
    "1. **Server Selection**: Takes `server_name` as a parameter (ArxivServer or WikipediaServer)\n",
    "2. **Session Management**: Uses `MCPClientSessionManager().open_session()` to connect\n",
    "3. **Toolkit Creation**: Converts MCP tools to AG2 format\n",
    "4. **Agent Creation**: Creates a temporary agent with access to the server's tools\n",
    "5. **Query Execution**: Runs the agent with the user's query\n",
    "6. **Result Extraction**: Returns the last message content\n",
    "7. **Cleanup**: Session automatically closes when exiting the context manager\n",
    "\n",
    "### Agent Handoffs\n",
    "\n",
    "The tool also configures a handoff back to the research assistant when the task is complete, creating a smooth workflow loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": [
    "TOOL_PROMPT = \"\"\"\n",
    "You are an MCP server tool.\n",
    "Your purpose is to identify the correct server to execute based on the user's query.\n",
    "\n",
    "Inputs:\n",
    "- query: (actual user query)\n",
    "- server_name: (name of the server to execute)\n",
    "\n",
    "You have two MCP servers to use:\n",
    "1. ArxivServer: to search for papers on arXiv\n",
    "2. WikipediaServer: to search for articles on Wikipedia\n",
    "\n",
    "NOTE:\n",
    "- Strictly return only the server name for server_name param (e.g., ArxivServer)\n",
    "- TERMINATE after response from the server\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "@tool(description=TOOL_PROMPT)\n",
    "async def run_mcp_agent_to_client(query: str, server_name: str) -> ReplyResult:\n",
    "    \"\"\"\n",
    "    Execute a query on the specified MCP server.\n",
    "\n",
    "    This tool:\n",
    "    1. Opens a session to the specified MCP server\n",
    "    2. Creates a toolkit from available MCP tools\n",
    "    3. Creates a temporary agent with those tools\n",
    "    4. Executes the query\n",
    "    5. Returns the result\n",
    "    \"\"\"\n",
    "    # Get the server configuration by name\n",
    "    server = get_server_config(mcp_config, server_name)\n",
    "\n",
    "    # Create a session manager and open a session\n",
    "    async with MCPClientSessionManager().open_session(server) as session:\n",
    "        # Initialize the session\n",
    "        await session.initialize()\n",
    "\n",
    "        # Get available tools from the server\n",
    "        agent_tool_prompt = await session.list_tools()\n",
    "\n",
    "        # Create a toolkit from the session\n",
    "        toolkit = await create_toolkit(session=session)\n",
    "\n",
    "        # Create a temporary agent for this server\n",
    "        agent = ConversableAgent(\n",
    "            name=\"mcp_agent\",\n",
    "            system_message=f\"You are an agent with access to {server_name} tools. Use them to answer queries.\",\n",
    "            llm_config=llm_config,\n",
    "            human_input_mode=\"NEVER\",\n",
    "        )\n",
    "\n",
    "        # Register MCP tools with the agent\n",
    "        toolkit.register_for_llm(agent)\n",
    "        toolkit.register_for_execution(agent)\n",
    "\n",
    "        # Configure handoff back to research assistant\n",
    "        agent.handoffs.add_llm_conditions([\n",
    "            OnCondition(\n",
    "                target=AgentTarget(research_assistant),\n",
    "                condition=StringLLMCondition(prompt=\"The research paper ids are fetched.\"),\n",
    "            ),\n",
    "        ])\n",
    "\n",
    "        # Execute the query using the MCP tools\n",
    "        result = await agent.a_run(\n",
    "            message=query + \" Use the following tools to answer the question: \" + str(agent_tool_prompt),\n",
    "            tools=toolkit.tools,\n",
    "            max_turns=5,\n",
    "        )\n",
    "\n",
    "        # Process results\n",
    "        res = await result.process()\n",
    "        last_message = await res.last_message()\n",
    "\n",
    "        # Return as ReplyResult with handoff to research assistant\n",
    "        return ReplyResult(message=str(last_message[\"content\"][-1]), target_agent=AgentTarget(research_assistant))\n",
    "\n",
    "\n",
    "print(\"Tool 'run_mcp_agent_to_client' defined successfully!\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "15",
   "metadata": {},
   "source": [
    "## Run the Research Assistant\n",
    "\n",
    "Now let's use the research assistant to query Wikipedia. The assistant will:\n",
    "1. Receive the user's query\n",
    "2. Call the `run_mcp_agent_to_client` tool with the appropriate server name\n",
    "3. The tool will open a session to WikipediaServer\n",
    "4. Execute the query using Wikipedia's MCP tools\n",
    "5. Return the results to the research assistant\n",
    "\n",
    "### Expected Workflow\n",
    "\n",
    "```\n",
    "User Query → Research Assistant → run_mcp_agent_to_client(query, \"WikipediaServer\")\n",
    "                ↑                                          ↓\n",
    "                └──────────── Results ←─────── MCP Agent (with Wikipedia tools)\n",
    "```\n",
    "\n",
    "> **Note**: Make sure the Wikipedia SSE server is running at `http://127.0.0.1:8000/sse` before executing this cell. You can start it with:\n",
    "> ```bash\n",
    "> python mcp/mcp_wikipedia.py sse --port 8000\n",
    "> ```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the research assistant with a query\n",
    "result = research_assistant.run(\n",
    "    message=\"Get me the latest news from Wikipedia using WikipediaServer\",\n",
    "    tools=[run_mcp_agent_to_client],\n",
    "    max_turns=2,\n",
    ").process()\n",
    "\n",
    "print(\"\\nResearch assistant completed!\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "17",
   "metadata": {},
   "source": [
    "## Alternative: Query ArXiv Papers\n",
    "\n",
    "Let's try querying the ArXiv server instead. This demonstrates the flexibility of the same tool working with different servers and transport protocols.\n",
    "\n",
    "The ArXiv server uses stdio transport (local process) while Wikipedia uses SSE (HTTP), but the `MCPClientSessionManager` handles both seamlessly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Query ArXiv for research papers\n",
    "result_arxiv = research_assistant.run(\n",
    "    message=\"Search for recent papers about large language models on ArxivServer\",\n",
    "    tools=[run_mcp_agent_to_client],\n",
    "    max_turns=2,\n",
    ").process()\n",
    "\n",
    "print(\"\\nArXiv query completed!\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "19",
   "metadata": {},
   "source": [
    "## Key Takeaways\n",
    "\n",
    "### Why Use MCPClientSessionManager This Way?\n",
    "\n",
    "This pattern of opening sessions inside tools provides several advantages:\n",
    "\n",
    "1. **Resource Efficiency**: Sessions are only opened when needed and automatically closed after use\n",
    "2. **Dynamic Server Selection**: The agent can choose which server to use based on the query\n",
    "3. **Isolation**: Each query gets a fresh session, preventing state pollution\n",
    "4. **Flexibility**: Easy to add new servers without changing the core workflow\n",
    "5. **Error Handling**: If one server fails, others remain unaffected\n",
    "\n",
    "### The Tool-Based Pattern\n",
    "\n",
    "By wrapping `MCPClientSessionManager` in a tool:\n",
    "- The LLM decides which server to use\n",
    "- Sessions are managed automatically\n",
    "- The same pattern works for any number of servers\n",
    "- Easy to extend with more server types\n",
    "\n",
    "### Best Practices\n",
    "\n",
    "1. **Use context managers**: Always use `async with` for proper cleanup\n",
    "2. **Initialize sessions**: Call `await session.initialize()` after opening\n",
    "3. **Unique server names**: Ensure each server has a distinct name in `MCPConfig`\n",
    "4. **Error handling**: Add try-except blocks for production use\n",
    "5. **Timeout configuration**: Set appropriate timeouts for SSE servers\n",
    "\n",
    "### Comparison with Traditional Approach\n",
    "\n",
    "**Traditional (keep session open):**\n",
    "```python\n",
    "async with stdio_client(...) as (read, write):\n",
    "    async with ClientSession(read, write) as session:\n",
    "        # Do all work here\n",
    "        # Session stays open entire time\n",
    "```\n",
    "\n",
    "**MCPClientSessionManager (on-demand):**\n",
    "```python\n",
    "async with MCPClientSessionManager().open_session(config) as session:\n",
    "    # Do specific work\n",
    "    # Session closes automatically\n",
    "```\n",
    "\n",
    "The `MCPClientSessionManager` pattern is ideal for:\n",
    "- Multi-server workflows\n",
    "- Long-running applications\n",
    "- Dynamic server selection\n",
    "- Resource-constrained environments\n",
    "\n",
    "### Next Steps\n",
    "\n",
    "- Add more MCP servers to `MCPConfig` (filesystem, database, APIs, etc.)\n",
    "- Implement intelligent routing logic in the research assistant\n",
    "- Add error handling and retry mechanisms\n",
    "- Create specialized agents for different server types\n",
    "- Build a server health check tool"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Learn how to use MCPClientSessionManager to dynamically manage multiple MCP server sessions (stdio and SSE) within tool functions for flexible, resource-efficient agent workflows.",
   "tags": [
    "MCP",
    "session management",
    "multi-server",
    "MCPClientSessionManager",
    "dynamic tools"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
